#!/usr/bin/env python3
"""
性能测试脚本 - 测试不同模型配置下的转录速度

测试目标：
1. tiny + CPU (int8)
2. base + CPU (int8)
3. small + CPU (int8) - 当前配置
4. medium + CPU (int8)

测试视频：10.2分钟的苏州留园视频
"""

import json
import time
import sys
from pathlib import Path
from datetime import datetime
from faster_whisper import WhisperModel

# 测试视频路径 - 使用路径代理
TEST_VIDEO = "/Volumes/Samsung_T5/数据爬取/B 站/关键词作品/狮子林/视频作品/BV13DGez8ETV苏州园林2-em class=keyword狮子林em 全网最详细游玩攻略.mp4"

VIDEO_DURATION = 509  # 秒 (8.5分钟)

# 测试配置
TEST_CONFIGS = [
    {
        "name": "tiny",
        "device": "cpu",
        "compute_type": "int8",
        "description": "最快速度，精度较低"
    },
    {
        "name": "base",
        "device": "cpu",
        "compute_type": "int8",
        "description": "快速，精度一般"
    },
    {
        "name": "small",
        "device": "cpu",
        "compute_type": "int8",
        "description": "平衡配置（当前使用）"
    },
    # medium 模型会花费较长时间，可选测试
    # {
    #     "name": "medium",
    #     "device": "cpu",
    #     "compute_type": "int8",
    #     "description": "高精度，较慢"
    # },
]


def format_time(seconds):
    """格式化时间显示"""
    if seconds < 60:
        return f"{seconds:.1f}秒"
    else:
        minutes = int(seconds // 60)
        secs = int(seconds % 60)
        return f"{minutes}分{secs}秒"


def calculate_speed_ratio(video_duration, transcribe_time):
    """计算倍速"""
    return video_duration / transcribe_time


def print_header():
    """打印测试标题"""
    print("\n" + "=" * 80)
    print(" " * 25 + "Whisper 转录性能测试")
    print("=" * 80)
    print(f"\n硬件信息:")
    print(f"  - CPU: Apple M4 Pro (12核)")
    print(f"  - RAM: 48 GB")
    print(f"  - 系统: macOS 15.6.1")
    print(f"\n测试视频:")
    print(f"  - 时长: {format_time(VIDEO_DURATION)} ({VIDEO_DURATION}秒)")
    print(f"  - 大小: 304 MB")
    print(f"  - 路径: .../{Path(TEST_VIDEO).name[:50]}...")
    print("\n" + "=" * 80 + "\n")


def test_model(config, test_number, total_tests):
    """测试单个模型配置"""
    model_name = config['name']
    device = config['device']
    compute_type = config['compute_type']
    description = config['description']

    print(f"\n{'─' * 80}")
    print(f"测试 {test_number}/{total_tests}: {model_name.upper()} 模型")
    print(f"{'─' * 80}")
    print(f"配置: {description}")
    print(f"  - 模型: {model_name}")
    print(f"  - 设备: {device}")
    print(f"  - 计算类型: {compute_type}")
    print()

    try:
        # 加载模型
        print(f"📦 加载模型中...", end="", flush=True)
        load_start = time.time()
        model = WhisperModel(
            model_name,
            device=device,
            compute_type=compute_type
        )
        load_time = time.time() - load_start
        print(f" 完成 ({load_time:.1f}秒)")

        # 执行转录
        print(f"🎬 开始转录 ({format_time(VIDEO_DURATION)}视频)...", flush=True)
        transcribe_start = time.time()

        segments, info = model.transcribe(
            TEST_VIDEO,
            language="zh",
            initial_prompt="以下是普通话的句子。"
        )

        # 收集结果
        segment_count = 0
        text_length = 0
        for segment in segments:
            segment_count += 1
            text_length += len(segment.text.strip())

            # 每10个片段显示一次进度
            if segment_count % 10 == 0:
                elapsed = time.time() - transcribe_start
                print(f"  进度: {segment_count} 片段 | 已耗时: {format_time(elapsed)}", flush=True)

        transcribe_time = time.time() - transcribe_start

        # 计算性能指标
        speed_ratio = calculate_speed_ratio(VIDEO_DURATION, transcribe_time)

        # 打印结果
        print(f"\n✅ 转录完成!")
        print(f"\n性能数据:")
        print(f"  ⏱️  转录耗时: {format_time(transcribe_time)} ({transcribe_time:.1f}秒)")
        print(f"  🚀 转录倍速: {speed_ratio:.2f}x")
        print(f"  📊 片段数量: {segment_count}")
        print(f"  📝 文本长度: {text_length} 字符")
        print(f"  🔤 识别语言: {info.language}")
        print(f"  ⚡ 处理速度: {VIDEO_DURATION/transcribe_time:.2f}秒视频/秒")

        if speed_ratio >= 10:
            emoji = "🔥🔥🔥"
        elif speed_ratio >= 5:
            emoji = "🔥🔥"
        elif speed_ratio >= 2:
            emoji = "🔥"
        else:
            emoji = "⚠️"

        print(f"\n  {emoji} 1小时视频预计需要: {format_time(3600/speed_ratio)}")

        return {
            "model": model_name,
            "device": device,
            "compute_type": compute_type,
            "load_time": load_time,
            "transcribe_time": transcribe_time,
            "speed_ratio": speed_ratio,
            "segment_count": segment_count,
            "text_length": text_length,
            "language": info.language,
            "success": True
        }

    except Exception as e:
        print(f"\n❌ 测试失败: {str(e)}")
        return {
            "model": model_name,
            "device": device,
            "compute_type": compute_type,
            "success": False,
            "error": str(e)
        }


def print_summary(results):
    """打印测试汇总"""
    print("\n" + "=" * 80)
    print(" " * 30 + "测试汇总")
    print("=" * 80)
    print()

    # 表头
    print(f"{'模型':<10} {'设备':<8} {'耗时':<15} {'倍速':<10} {'1小时视频预计':<15}")
    print("─" * 80)

    # 成功的测试结果
    successful_results = [r for r in results if r['success']]

    for result in successful_results:
        model = result['model']
        device = result['device']
        time_str = format_time(result['transcribe_time'])
        speed = f"{result['speed_ratio']:.2f}x"
        hour_estimate = format_time(3600 / result['speed_ratio'])

        print(f"{model:<10} {device:<8} {time_str:<15} {speed:<10} {hour_estimate:<15}")

    # 失败的测试
    failed_results = [r for r in results if not r['success']]
    if failed_results:
        print("\n失败的测试:")
        for result in failed_results:
            print(f"  ❌ {result['model']} - {result['error']}")

    # 推荐配置
    if successful_results:
        print("\n" + "=" * 80)
        print("📊 推荐配置:")
        print()

        fastest = max(successful_results, key=lambda x: x['speed_ratio'])
        print(f"🏆 最快速度: {fastest['model']} ({fastest['speed_ratio']:.2f}x)")
        print(f"   1小时视频仅需: {format_time(3600/fastest['speed_ratio'])}")

        if fastest['model'] == 'tiny':
            print("\n💡 提示: tiny 模型速度最快，但精度较低")
            print("   如果需要更高精度，建议使用 small 模型")
        elif fastest['model'] == 'small':
            print("\n💡 这是速度和精度的最佳平衡点！")

    print("\n" + "=" * 80)


def save_results(results):
    """保存测试结果到JSON文件"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"benchmark_results_{timestamp}.json"

    data = {
        "timestamp": datetime.now().isoformat(),
        "hardware": {
            "cpu": "Apple M4 Pro",
            "cores": 12,
            "ram": "48 GB",
            "os": "macOS 15.6.1"
        },
        "test_video": {
            "duration": VIDEO_DURATION,
            "size": "304 MB"
        },
        "results": results
    }

    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    print(f"\n💾 结果已保存到: {filename}")


def main():
    """主测试流程"""
    print_header()

    # 检查测试视频是否存在
    if not Path(TEST_VIDEO).exists():
        print(f"❌ 错误: 测试视频不存在")
        print(f"路径: {TEST_VIDEO}")
        sys.exit(1)

    # 确认开始测试
    print(f"将测试 {len(TEST_CONFIGS)} 个配置，预计总耗时: 15-30分钟")
    print("\n⚠️  测试期间会占用较多CPU资源，建议不要运行其他重度任务")

    # 支持命令行参数 --yes 跳过确认
    if '--yes' not in sys.argv and '-y' not in sys.argv:
        try:
            response = input("\n是否开始测试? (y/n): ").strip().lower()
            if response != 'y':
                print("已取消测试")
                sys.exit(0)
        except EOFError:
            print("\n自动开始测试...")
    else:
        print("\n自动开始测试...")

    # 执行测试
    results = []
    total_start = time.time()

    for i, config in enumerate(TEST_CONFIGS, 1):
        result = test_model(config, i, len(TEST_CONFIGS))
        results.append(result)

        # 测试之间稍作休息
        if i < len(TEST_CONFIGS):
            print("\n⏸️  休息 5 秒...\n")
            time.sleep(5)

    total_time = time.time() - total_start

    # 打印汇总
    print_summary(results)

    print(f"\n⏱️  总测试时长: {format_time(total_time)}")

    # 保存结果
    save_results(results)

    print("\n✨ 所有测试完成！")


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠️  测试被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ 发生错误: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
